{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  ContML.CPSTrafo
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- CPS transformation.
--
--------------------------------------------------------------------------------


module ContML.CPSTrafo
  ( 

    cps

  ) where


import ContML.CPS
import ContML.FreshMon
import qualified ContML.Syntax as S
import ContML.Types
import ContML.VarId


type CVar = VarId

cps :: S.Expr -> Mon Term
cps e0 = fresh 'h' TyUnknown >>= \h -> cpsT e0 h ctx
  where
    ctx x = do { k <- fresh 'k' TyUnknown; return $ AppCont k x }



cpsT :: S.Expr -> CVar -> (VarId -> Mon Term) -> Mon Term
cpsT (S.Int i)                      _ kT = do
    x     <- fresh 'x' TyInt
    ans   <- kT x
    return $ LetVal x (Int i) ans

cpsT (S.Var x)                      _ kT = kT x

cpsT (S.App e1 e2)                  h kT = do
    k     <- fresh 'k'  TyUnknown
    x     <- fresh 'x' TyUnknown
    ans   <- kT x
    cpsT e1 h $ \x1 -> 
        cpsT e2 h $ \x2 -> 
            return $ LetCont [ContDef k x ans] (App x1 k h x2)

cpsT (S.Fun x e)                    _ kT = do 
    f     <- fresh 'f' TyUnknown
    h'    <- fresh 'k'  TyUnknown
    k     <- fresh 'k'  TyUnknown
    eT    <- cpsTail e h' k
    ans   <- kT f
    return $ LetFun [FunDef f k h' x eT] ans

cpsT (S.Pair e1 e2)                 h kT = do
    x     <- fresh 'x' (TyPair TyUnknown TyUnknown)
    ans   <- kT x
    cpsT e1 h $ \x1 -> 
        cpsT e2 h $ \x2 -> 
            return $ LetVal x (Pair x1 x2) ans

cpsT (S.Tag i e)                    h kT = do
    x     <- fresh 'x' TyUnknown
    ans   <- kT x
    cpsT e h $ \z -> 
        return $ LetVal x (Tag i z) ans

cpsT (S.Unit)                       _ kT = do
    x     <- fresh 'x' TyUnit
    ans   <- kT x
    return $ LetVal x Unit ans

cpsT (S.Unpair i e)                 h kT = do
    x     <- fresh 'x' TyUnknown
    ans   <- kT x
    cpsT e h $ \z -> 
        return $ LetProj x i z ans

cpsT (S.LetVal x e1 e2)             h kT = do
    j     <- fresh 'j' TyUnknown
    body  <- cpsT e2 h kT
    ans   <- cpsTail e1 h j
    return $ LetCont [ContDef j x body] ans

cpsT (S.LetFun ds e)                h kT = do
    ds'   <- mapM cpsFunDef ds
    ans   <- cpsT e h kT
    return $ LetFun ds' ans

cpsT (S.Raise e)                    h _  = do
    cpsT e h $ \z -> 
        return $ AppCont h z

cpsT (S.Handle e1 x e2)             h kT = do
    j     <- fresh 'j' TyUnknown
    h'    <- fresh 'h' TyUnknown
    body  <- cpsTail e2 h j
    final <- cpsTail e1 h' j
    ans   <- kT x
    return $ LetCont [ContDef j x ans] (LetCont [ContDef h' x body] final)

cpsT (S.Case e x1 e1 x2 e2)         h kT = do
    x     <- fresh 'x' TyUnknown
    k1    <- fresh 'k' TyUnknown
    k2    <- fresh 'k' TyUnknown     
    j     <- fresh 'j' TyUnknown
    e1'   <- cpsTail e1 h j
    e2'   <- cpsTail e2 h j
    ans   <- kT x
    cpsT e h $ \z -> 
        return $ LetCont [ContDef j x ans] 
                         (LetCont [ContDef k1 x1 e1']
                                  (LetCont [ContDef k2 x2 e2']
                                           (Case z k1 k2)))
  

cpsFunDef :: S.Def -> Mon FunDef
cpsFunDef (S.Def f x e)                  = do
    k   <- fresh 'k' TyUnknown
    h   <- fresh 'k' TyUnknown
    ans <- cpsTail e h k 
    return $ FunDef f k h x ans


-- | Aka:
--
-- > (|.|)
--
cpsTail :: S.Expr -> CVar -> CVar -> Mon Term
cpsTail (S.Int i)                    _ k = do
    x     <- fresh 'x' TyInt
    return $ LetVal x (Int i) (AppCont k x)

cpsTail (S.Var x)                   _ k = return $ AppCont k x

cpsTail (S.App e1 e2)               h k = 
    cpsT e1 h $ \x1 -> 
        cpsT e2 h $ \x2 -> 
            return $ App x1 k h x2

cpsTail (S.Fun x e)                 h k = do
    f     <- fresh 'f' (TyFun TyUnknown TyUnknown)
    j     <- fresh 'k' TyUnknown
    e'    <- cpsTail e h j
    return $ LetVal f (Abs j x e')  (AppCont k f)


cpsTail (S.Pair e1 e2)              h k = do
    x     <- fresh 'x' TyUnknown
    cpsT e1 h $ \x1 -> 
        cpsT e2 h $ \x2 -> 
            return $ LetVal x (Pair x1 x2) (AppCont k x)

cpsTail (S.Tag i e)                 h k = do
    x     <- fresh 'x' TyUnknown
    cpsT e h $ \z -> 
        return $ LetVal x (Tag i z) (AppCont k x)

cpsTail (S.Unit)                    _ k = do
    x     <- fresh 'x' TyUnit
    return $ LetVal x Unit (AppCont k x)

cpsTail (S.Unpair i e)              h k = do
    x     <- fresh 'x' TyUnknown
    cpsT e h $ \z -> 
        return $ LetProj x i z (AppCont k x)

cpsTail (S.LetVal x e1 e2)          h k = do
    j     <- fresh 'j' TyUnknown
    e1'   <- cpsTail e1 h j
    e2'   <- cpsTail e2 h k
    return $ LetCont [ContDef j x e2'] e1'

cpsTail (S.LetFun ds e)             h k = do
    ds'   <- mapM cpsFunDef ds
    e'    <- cpsTail e h k
    return $ LetFun ds' e'

cpsTail (S.Raise e)                 h _ = do
    cpsT e h $ \z -> 
        return $ AppCont h z

cpsTail (S.Handle e1 x e2)          h k = do
    h'    <- fresh 'h' TyUnknown
    e1'   <- cpsTail e1 h' k
    e2'   <- cpsTail e2 h k
    return $ LetCont [ContDef h' x e2'] e1'

cpsTail (S.Case e x1 e1 x2 e2)      h k = do
    k1    <- fresh 'k' TyUnknown
    k2    <- fresh 'k' TyUnknown     
    e1'   <- cpsTail e1 h k
    e2'   <- cpsTail e2 h k
    cpsT e h $ \z -> 
        return $ LetCont [ContDef k1 x1 e1'] 
                         (LetCont [ContDef k2 x2 e2']
                                  (Case z k1 k2))